#!/usr/env/bin python
# -*- coding: utf-8 -*-

# @Time    : 2018/8/21 18:30
# @Author  : yangdy
# @Email   : yangdy@egu360.com
# @File    : 中级-每个节点的右向指针.py
# @Software: PyCharm

# Definition for binary tree with next pointer.


class TreeLinkNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None
        self.next = None


class Solution:
    # @param root, a tree link node
    # @return nothing
    def connect(self, root):
        que = []
        if root:
            que.append(root)
        while que:
            q = []
            lq = len(que)
            for i in range(lq):
                que[i].next = que[i + 1] if i + 1 < lq else None
                if que[i].left:
                    q.append(que[i].left)
                if que[i].right:
                    q.append(que[i].right)
            que = q


def display(nd, dp=0):
    if not nd:
        return
    print('  '*dp, nd.val, '(', nd.next.val if nd.next else None, ')')
    display(nd.left, dp+1)
    display(nd.right, dp+1)


def main():
    solve = Solution()
    # {1, 2, 3}
    root = TreeLinkNode(1)
    root.left = TreeLinkNode(2)
    root.right = TreeLinkNode(3)
    display(root)
    solve.connect(root)
    display(root)


if __name__ == '__main__':
    main()
